import os
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('-c', '--cuda', type=str, default='2', help='CUDA device id')
args, unknown = parser.parse_known_args()
os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda

import gradio as gr
import torch
from transformers import AutoConfig, AutoModelForCausalLM
from janus.models import MultiModalityCausalLM, VLChatProcessor
from janus.utils.io import load_pil_images
from PIL import Image

import numpy as np
import os
import time
from util_for_os import ose,osj
# import spaces  # Import spaces for ZeroGPU compatibility


# Load model and processor
# model_path = "deepseek-ai/Janus-Pro-7B"
model_path = "/mnt/nas/zhangshu/Janus-Pro-7B"
config = AutoConfig.from_pretrained(model_path)
language_config = config.language_config
language_config._attn_implementation = 'eager'
vl_gpt = AutoModelForCausalLM.from_pretrained(model_path,
                                             language_config=language_config,
                                             trust_remote_code=True)
if torch.cuda.is_available():
    vl_gpt = vl_gpt.to(torch.bfloat16).cuda()
else:
    vl_gpt = vl_gpt.to(torch.float16)

vl_chat_processor = VLChatProcessor.from_pretrained(model_path)
tokenizer = vl_chat_processor.tokenizer
cuda_device = 'cuda' if torch.cuda.is_available() else 'cpu'


@torch.inference_mode()
# @spaces.GPU(duration=120) 
# Multimodal Understanding function
def multimodal_understanding(image, question, seed=42, top_p=0.95, temperature=0.1):
    # Clear CUDA cache before generating
    torch.cuda.empty_cache()
    
    # set seed
    torch.manual_seed(seed)
    np.random.seed(seed)
    torch.cuda.manual_seed(seed)
    
    conversation = [
        {
            "role": "<|User|>",
            "content": f"<image_placeholder>\n{question}",
            "images": [image],
        },
        {"role": "<|Assistant|>", "content": ""},
    ]
    
    pil_images = [Image.fromarray(image)]
    prepare_inputs = vl_chat_processor(
        conversations=conversation, images=pil_images, force_batchify=True
    ).to(cuda_device, dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float16)
    
    
    inputs_embeds = vl_gpt.prepare_inputs_embeds(**prepare_inputs)
    
    outputs = vl_gpt.language_model.generate(
        inputs_embeds=inputs_embeds,
        attention_mask=prepare_inputs.attention_mask,
        pad_token_id=tokenizer.eos_token_id,
        bos_token_id=tokenizer.bos_token_id,
        eos_token_id=tokenizer.eos_token_id,
        max_new_tokens=512,
        do_sample=False if temperature == 0 else True,
        use_cache=True,
        temperature=temperature,
        top_p=top_p,
    )
    
    answer = tokenizer.decode(outputs[0].cpu().tolist(), skip_special_tokens=True)
    return answer

from os.path import join as osj



dates = ['1106']
categories = ['jacket','hoodie','formalatt','weddress']

# dates = ['0917']
# categories = ['slipdress']



for date in dates:
    for category in categories:
        dir = f'/mnt/nas/datasets/diction/{category}{date}_img_clo_diff'
        
        assert ose( dir )
        assert ose( (osj( dir , 'names.txt' )) )
    
for date in dates:
    for category in categories:
        dir = f'/mnt/nas/datasets/diction/{category}{date}_img_clo_diff'
        

        save_clothing_dir = f'{dir}_furry'
        save_human_dir = f'{dir}_furryless'
        import os,shutil,cv2
        # filenames = os.listdir(dir)
        if os.path.exists(save_clothing_dir):shutil.rmtree(save_clothing_dir)
        os.makedirs(save_clothing_dir)
        if os.path.exists(save_human_dir):shutil.rmtree(save_human_dir)
        os.makedirs(save_human_dir)

        Q1 = '如何描述这件衣服'  # no
        # Q2 = '是否是一件完整的衣服'   # yes
        # Q3 = 'Is the background white? Just answer yes or no.'   # yes
        # Q4 = '是否是鞋子' # no
        # Q5 = '是否是衣服' # yes

        keywords = ['毛茸茸','毛绒','长毛',
                    '毛皮','羊羔毛']

        # check_no = lambda ans : 'no' in ans.lower() or '否' in ans or '不是' in ans
        check_yes = lambda ans : any( [ k in ans for k in keywords] )

        # count = 0
        with open(  osj( dir , 'names.txt' ) ,encoding='utf-8') as f:
            names = f.readlines() 
        import pdb
        from tqdm import tqdm
        # pdb.set_trace()
        # for entry in os.scandir( dir ):
        get_answer = lambda filepath,q : multimodal_understanding( cv2.imread( filepath ) , q )
        for filename in tqdm(names):
            # if entry.is_file() and entry.name.endswith('.jpg'):
            filename = filename.strip()
            if filename.endswith('.jpg'):
                # filename = entry.name
                pass
            else:continue

            # count += 1
            # print('\rprocess num : ',count,end='',flush=True)    
            
            filepath = os.path.join( dir , filename)


            # results = model.predict(filepath)
            
            # A1 = multimodal_understanding( cv2.imread( filepath ) , Q1 )
            # A2 = multimodal_understanding( cv2.imread( filepath ) , Q2 )
            # A3 = multimodal_understanding( cv2.imread( filepath ) , Q3 )

            ans = get_answer( filepath , Q1 )
            
            # pdb.set_trace()
            
            if check_yes( ans ) :
                # check_yes(  get_answer( filepath , Q2 )  ) and \
                #     check_yes(  get_answer( filepath , Q3 ) ) and \
                #         check_no( get_answer(filepath , Q4) )  and \
                #             check_yes( get_answer( filepath , Q5 ) ):
                print(osj(save_clothing_dir,filename), ans)
                shutil.copy2( osj(dir , filename) , osj(save_clothing_dir,filename) )
            else:
                shutil.copy2( osj(dir , filename) , osj(save_human_dir,filename) )
        
'''
/mnt/nas/datasets/diction/images/               筛选后的图(最终存放位置)
/mnt/nas/datasets/diction/ZipArchive            不筛选的zip
/mnt/nas/datasets/diction/ZipArchive07XX        解压后的图片
/mnt/nas/datasets/diction/ZipArchive07XX_clo    Janus筛选后的clo图 (包含重复图,可能出现几张不完整或者背景不干净,错误概率比例小于2%)
/mnt/nas/datasets/diction/ZipArchive07XX_clo_sim     Janus筛选后的clo图 (包含重复图 clo)
/mnt/nas/datasets/diction/ZipArchive07XX_clo_diff    Janus筛选后的clo图 (包含不重复图 clo,可能出现几张不完整或者背景不干净,错误概率比例小于2%)
/mnt/nas/datasets/diction/ZipArchive07XX_human      Janus筛选后的非clo图 (人 / 不完整)

date=1020
bash restore.sh /mnt/nas/datasets/diction/dress"$date"
bash restore.sh /mnt/nas/datasets/diction/suit"$date"
bash restore.sh /mnt/nas/datasets/diction/suitset"$date"
bash restore.sh /mnt/nas/datasets/diction/trousers"$date"

date=1020
bash restore.sh /mnt/nas/datasets/diction/dress"$date"_img
bash restore.sh /mnt/nas/datasets/diction/suit"$date"_img
bash restore.sh /mnt/nas/datasets/diction/suitset"$date"_img
bash restore.sh /mnt/nas/datasets/diction/trousers"$date"_img

date=1020
bash restore.sh /mnt/nas/datasets/diction/dress"$date"_img_clo
bash restore.sh /mnt/nas/datasets/diction/suit"$date"_img_clo
bash restore.sh /mnt/nas/datasets/diction/suitset"$date"_img_clo
bash restore.sh /mnt/nas/datasets/diction/trousers"$date"_img_clo

date=1106
bash restore.sh /mnt/nas/datasets/diction/dress"$date"_img_clo_diff
bash restore.sh /mnt/nas/datasets/diction/suit"$date"_img_clo_diff
bash restore.sh /mnt/nas/datasets/diction/suitset"$date"_img_clo_diff
bash restore.sh /mnt/nas/datasets/diction/trousers"$date"_img_clo_diff


'''